package goYmodem

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"goYmodem/crc16"
	"goYmodem/utils"
	"io"
	"io/ioutil"
	"strconv"
	"time"
)

const (
	YMODEM_CODE_NONE byte = 0x0
	YMODEM_CODE_SOH  byte = 0x01  /* start of 128-byte data packet */
	YMODEM_CODE_STX  byte = 0x02  /* start of 1024-byte data packet */
	YMODEM_CODE_EOT  byte = 0x04  /* end of transmission */
	YMODEM_CODE_ACK  byte = 0x06  /* receive OK */
	YMODEM_CODE_NAK  byte = 0x015 /* receiver error; retry */
	YMODEM_CODE_CAN  byte = 0x018 /* two of these in succession aborts transfer */
	YMODEM_CODE_CRC  byte = 0x43  /* use in place of first NAK for CRC mode */
)

const (
	YMODEM_STAGE_NONE = iota
	/* set when C is send */
	YMODEM_STAGE_ESTABLISHING
	/* set when we've got the packet 0 and sent ACK and second C */
	YMODEM_STAGE_ESTABLISHED
	/* set when the sender respond to our second C and recviever got a real
	 * data packet. */
	YMODEM_STAGE_TRANSMITTING
	/* set when the sender send a EOT */
	YMODEM_STAGE_FINISHING
	/* set when transmission is really finished, i.e., after the NAK, C, final
	 * NULL packet stuff. */
	YMODEM_STAGE_FINISHED
)

/**
ymodem 协议帧格式
帧头-----SOH 128字节 STX 1024 字节 ， 1个字节
包号------------------------------- 1个字节
包号反码---------------------------- 1个字节
信息块------------------------------133 字节 或者1024 字节
CRC--------------------------------2个字节，高位在前地位在后
*/

type Ymodem struct {
	FileName string
	FileSize uint32
	data     []byte
	stage    int

	BeginPack func([]byte)
	DataPack  func([]byte)
	EndPack   func([]byte)
	io.ReadWriter
}

func (y *Ymodem) Send(data []byte) {
	crcNak := 1
	var ch byte
	retry := 0
	//获取握手信号
	for ; retry < 60 && ch != YMODEM_CODE_CRC; retry++ {
		fmt.Println("Try to get C cmd: ", retry)
		time.Sleep(time.Millisecond * 100)
		err := binary.Read(y, binary.LittleEndian, &ch)
		if err != nil {
			continue
		}
	}
	//接收到接收方发送的 C
	if ch == YMODEM_CODE_CRC {
		for {
			//发送第一帧数据包
			y.SendPack0()
			binary.Read(y, binary.LittleEndian, &ch)

			if ch == YMODEM_CODE_ACK {
				binary.Read(y, binary.LittleEndian, &ch)
				if ch == YMODEM_CODE_CRC {
					y.SendDataPackets(data)
					return
				} else if ch == YMODEM_CODE_CRC && crcNak > 0 {
					crcNak = 0
					continue
				} else if ch != YMODEM_CODE_NAK || crcNak > 0 {
					break
				}
			}
		}
	} else {
		binary.Write(y, binary.LittleEndian, YMODEM_CODE_CAN)
		binary.Write(y, binary.LittleEndian, YMODEM_CODE_CAN)
	}
}

func (y *Ymodem) SendDataPackets(data []byte) {
	size := y.FileSize
	retry := 0
	var blockNo uint8 = 1
	var sendSize uint32 = 0
	var sendNum uint32
	var ch byte
	dataRead := bytes.NewReader(data)
	for size > 0 {
		if size > 128 {
			sendSize = 1024
		} else {
			sendSize = 128
		}
		if sendSize > 1024 {
			sendSize = 1024
		}
		var bufSize uint32
		if size < sendSize {
			bufSize = size
		} else {
			bufSize = sendSize
		}
		dataBuf := make([]byte, bufSize)
		binary.Read(dataRead, binary.LittleEndian, &dataBuf)
		for len(dataBuf) < int(sendSize) {
			dataBuf = append(dataBuf, 0x1a)
		}
		y.SendPacket(dataBuf, blockNo)
		binary.Read(y, binary.LittleEndian, &ch)
		//发送数据的钩子函数
		switch ch {
		case YMODEM_CODE_ACK:
			retry = 0
			blockNo++
			sendNum += sendSize
			if size < sendSize {
				size = 0
			} else {
				size -= sendSize
			}
		case YMODEM_CODE_NAK:
			if retry++; retry > 5 {
				return
			}
		default:
			return
		}
	}
	//数据发送完毕
	binary.Write(y, binary.LittleEndian, YMODEM_CODE_EOT)
	binary.Read(y, binary.LittleEndian, &ch)
	for ch != YMODEM_CODE_ACK {
		binary.Write(y, binary.LittleEndian, YMODEM_CODE_EOT)
		binary.Read(y, binary.LittleEndian, &ch)
	}
	if ch == YMODEM_CODE_ACK {
		binary.Read(y, binary.LittleEndian, &ch)
		if ch == YMODEM_CODE_CRC {
			for ch != YMODEM_CODE_ACK {
				y.SendPackEnd()
				binary.Read(y, binary.LittleEndian, &ch)
			}
		}
	}
	//y.EndPack(nil)
}

func (y *Ymodem) SendFile(FileName string) {
	pathExist, _ := utils.PathExists(FileName)
	if pathExist != true {
		fmt.Println("file is not exist ", FileName)
		return
	}
	data, err := ioutil.ReadFile(FileName)
	if err != nil {
		fmt.Println(err)
		return
	}
	y.FileName = FileName
	y.FileSize = uint32(len(data))
	y.Send(data)
}

func (y *Ymodem) SendPack0() {
	if y.FileName == "" {
		return
	}
	dataBuff := bytes.NewBuffer([]byte{})
	if err := binary.Write(dataBuff, binary.LittleEndian, y.FileName); err != nil {
		return
	}
	binary.Write(dataBuff, binary.LittleEndian, byte(0))
	if err := binary.Write(dataBuff, binary.LittleEndian, utils.U32toStr(y.FileSize)); err != nil {
		return
	}
	binary.Write(dataBuff, binary.LittleEndian, byte(0))
	zero := make([]byte, 128-dataBuff.Len())
	binary.Write(dataBuff, binary.LittleEndian, zero)
	//发送第一帧
	y.SendPacket(dataBuff.Bytes(), 0)
}

func (y *Ymodem) SendPackEnd() {

	block := make([]byte, 128)
	//发送第一帧
	y.SendPacket(block, 0)
}

func (y *Ymodem) SendPacket(data []byte, blockNo uint8) {
	crc := crc16.CheckSum(data)
	var packet [3]byte
	if len(data) == 128 {
		packet[0] = YMODEM_CODE_SOH
	} else {
		packet[0] = YMODEM_CODE_STX
	}
	//发送头部，确定消息大小，包号
	packet[1] = blockNo & 0xff
	//发送消息编号，包号反码
	packet[2] = (^blockNo) & 0xff
	binary.Write(y, binary.LittleEndian, packet)
	binary.Write(y, binary.LittleEndian, data)
	binary.Write(y, binary.BigEndian, crc)
}

func (y *Ymodem) ReceivePackage() int {
	var c byte

	binary.Read(y, binary.LittleEndian, &c)
	if c < 0 {
		return -1
	}
	var packSize uint32
	switch c {
	case YMODEM_CODE_SOH:
		packSize = 128
	case YMODEM_CODE_STX:
		packSize = 1024
	case YMODEM_CODE_EOT:
		packSize = 0
	case YMODEM_CODE_CAN:
		binary.Read(y, binary.LittleEndian, &c)
		if c == YMODEM_CODE_CAN {
			return 0
		}
	default:
		return -1
	}
	y.data = []byte{}
	y.data = append(y.data, c)

	var buf = make([]byte, packSize+4)
	binary.Read(y, binary.LittleEndian, &buf)

	y.data = append(y.data, buf...)

	if packSize != 0 {
		crcData := make([]byte, packSize)
		read := bytes.NewReader(y.data)
		read.Seek(3, io.SeekCurrent)
		binary.Read(read, binary.LittleEndian, &crcData)
		var crc uint16
		binary.Read(read, binary.LittleEndian, &crc)
		if crc == crc16.CheckSum(crcData) {
			return 0
		} else {
			return -1
		}
	}
	return 0
}

func (y *Ymodem) Receive() int {
	ret := -1
	y.stage = YMODEM_STAGE_NONE
	ret = y.Handshake()
	if ret != 0 {
		return ret
	}
	for {
		ret = y.ReceiveFileData()
		if ret != 0 {
			return ret
		}
		ret = y.ReceiveFinish()
		if ret != 0 {
			return ret
		}
		if y.stage == YMODEM_STAGE_FINISHED {
			break
		}
	}
	return 0
}

func (y *Ymodem) Handshake() int {
	var timeout uint32 = 60
	y.stage = YMODEM_STAGE_FINISHING
	for ; timeout <= 0; timeout-- {
		binary.Write(y, binary.LittleEndian, YMODEM_CODE_CRC)
		ret := y.ReceivePackage()
		if ret == 0 && len(y.data) > 128 {
			reader := bytes.NewReader(y.data)
			reader.Seek(3, io.SeekCurrent)
			y.FileName = utils.ReadString(reader)
			sizebuf := utils.ReadString(reader)
			size, _ := strconv.Atoi(sizebuf)
			y.FileSize = uint32(size)
			binary.Write(y, binary.LittleEndian, YMODEM_CODE_ACK)
			binary.Write(y, binary.LittleEndian, YMODEM_CODE_CRC)
			break
		}
		if timeout <= 0 {
			return -1
		}
	}
	return 0
}

func (y *Ymodem) ReceiveFileData() int {
	retry := 0
	y.stage = YMODEM_STAGE_TRANSMITTING
	for {
		ret := y.ReceivePackage()
		if ret == 0 {
			if len(y.data) > 0 {
				if y.DataPack != nil {
					y.DataPack(y.data)
				}
				binary.Write(y, binary.LittleEndian, YMODEM_CODE_ACK)
			} else if y.data[0] == YMODEM_CODE_CAN {
				for i := 0; i < 7; i++ {
					binary.Write(y, binary.LittleEndian, YMODEM_CODE_CAN)
				}
			} else if y.data[0] == YMODEM_CODE_EOT {
				return 0
			}
		} else {
			if retry++; retry > 60 {
				break
			} else {
			}
		}
	}
	return -1
}

func (y *Ymodem) ReceiveFinish() int {
	y.stage = YMODEM_STAGE_FINISHING
	binary.Write(y, binary.LittleEndian, YMODEM_CODE_NAK)

	var ch byte
	binary.Read(y, binary.LittleEndian, &ch)
	if ch != YMODEM_CODE_EOT {
		return -1
	}
	binary.Write(y, binary.LittleEndian, YMODEM_CODE_ACK)
	binary.Write(y, binary.LittleEndian, YMODEM_CODE_CRC)

	ret := y.ReceivePackage()
	if ret != 0 {
		return -1
	}
	if len(y.data) != 128 {
		return -1
	}
	binary.Write(y, binary.LittleEndian, YMODEM_CODE_ACK)
	//y.EndPack(y.data)

	if y.data[3] != 0 {
		//y.BeginPack(y.data)
		binary.Write(y, binary.LittleEndian, YMODEM_CODE_CRC)
		return 0
	}
	y.stage = YMODEM_STAGE_FINISHED
	return 0
}
